import os
import sys
from datetime import datetime
from config.config_manager import ConfigManager
from db.mysql_db import MySQLDB
from db.kingbase_db import KingbaseDB
from migration.metadata_migrator import MetadataMigrator
from migration.data_migrator import DataMigrator
from task.task_tracker import TaskTracker
from utils.logger import Logger
from utils.common import CommonUtils

class MySQL2KingbaseESMigrator:
    def __init__(self, task_name):
        self.task_name = task_name
        
        # 初始化日志
        self.logger = Logger(task_name)
        
        # 初始化配置管理器
        self.config_manager = ConfigManager(task_name)
        
        # 初始化数据库连接
        self.mysql_db = None
        self.kingbase_db = None
        
        # 初始化迁移器
        self.metadata_migrator = None
        self.data_migrator = None
        
        # 初始化任务跟踪器
        self.task_tracker = TaskTracker(task_name, self.logger)
    
    def initialize(self):
        """
        初始化所有组件
        """
        try:
            # 加载配置
            self.logger.info(f"开始加载配置: {self.task_name}")
            config = self.config_manager.load_config()
            
            # 初始化MySQL连接
            self.logger.info(f"开始连接MySQL数据库: {config['source_db']['host']}:{config['source_db']['port']}")
            self.mysql_db = MySQLDB(config['source_db'])
            self.mysql_db.connect()
            
            # 初始化KingbaseES连接
            self.logger.info(f"开始连接KingbaseES数据库: {config['target_db']['host']}:{config['target_db']['port']}")
            self.kingbase_db = KingbaseDB(config['target_db'])
            self.kingbase_db.connect()
            
            # 初始化元数据迁移器
            self.metadata_migrator = MetadataMigrator(
                self.mysql_db,
                self.kingbase_db,
                {
                    'target_schema': config['target_db']['schema'],
                    'skip_existing_tables': config['migration']['skip_existing_tables'],
                    'include_tables': config['migration']['include_tables'],
                    'exclude_tables': config['migration']['exclude_tables'],
                    'export_view_sql': config['migration']['export_view_sql']
                },
                self.logger,
                self.task_name
            )
            
            # 初始化数据迁移器
            self.data_migrator = DataMigrator(
                self.mysql_db,
                self.kingbase_db,
                {
                    'target_schema': config['target_db']['schema'],
                    'batch_size': config['migration']['batch_size']
                },
                self.logger,
                self.task_name
            )
            
            self.logger.info("初始化完成")
            return True
        except Exception as e:
            self.logger.error(f"初始化失败: {e}")
            self.cleanup()
            return False
    
    def run(self):
        """
        运行迁移流程
        """
        try:
            # 初始化
            if not self.initialize():
                return False
            
            # 开始任务
            self.task_tracker.start_task()
            
            # 获取续传信息
            resume_info = self.task_tracker.get_resume_info()
            
            # 迁移元数据（表结构）
            tables, views = self.metadata_migrator.migrate_metadata()
            
            # 更新任务跟踪器
            self.task_tracker.set_tables(tables)
            self.task_tracker.set_views(views)
            
            # 获取待迁移的表（包括续传的表）
            pending_tables = self.task_tracker.get_pending_tables()
            
            if pending_tables:
                self.logger.info(f"开始迁移数据，共 {len(pending_tables)} 个表需要迁移")
                
                # 迁移数据
                migration_results = self.data_migrator.migrate_data(pending_tables, resume_info)
                
                # 更新任务跟踪器
                self.task_tracker.set_migration_results(migration_results)
            else:
                self.logger.info("所有表都已迁移完成，无需重新迁移")
            
            # 生成迁移报告
            self.task_tracker.generate_report()
            
            # 完成任务
            self.task_tracker.finish_task(success=True)
            
            self.logger.info("迁移任务已成功完成")
            return True
        except Exception as e:
            self.logger.error(f"迁移任务执行失败: {e}")
            self.task_tracker.finish_task(success=False, error=str(e))
            return False
        finally:
            self.cleanup()
    
    def cleanup(self):
        """
        清理资源
        """
        try:
            # 关闭数据库连接
            if self.mysql_db:
                self.mysql_db.disconnect()
                self.logger.info("MySQL数据库连接已关闭")
            
            if self.kingbase_db:
                self.kingbase_db.disconnect()
                self.logger.info("KingbaseES数据库连接已关闭")
        except Exception as e:
            self.logger.error(f"清理资源时发生错误: {e}")

def main():
    """
    主函数
    """
    # 检查命令行参数
    if len(sys.argv) < 2:
        print("用法: python main.py <任务名称>")
        sys.exit(1)
    
    # 获取任务名称
    task_name = sys.argv[1]
    
    # 创建迁移器实例并运行
    migrator = MySQL2KingbaseESMigrator(task_name)
    success = migrator.run()
    
    # 根据执行结果退出
    sys.exit(0 if success else 1)

if __name__ == "__main__":
    main()